{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06e5af38-85ad-46b6-a5ea-52a7ff947853",
   "metadata": {},
   "source": [
    "# Labelme转YOLO-Keypoint-单个文件\n",
    "\n",
    "同济子豪兄 2023-4-16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a8a2a1-fb46-4d51-8b4f-cd0568bc13f8",
   "metadata": {},
   "source": [
    "## 导入工具包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5969d148-19e6-4b2e-9797-01105b87b56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31ef8f72-7733-4d84-8250-c48583f416a3",
   "metadata": {},
   "source": [
    "## 数据集类别信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bc933b72-ef9b-471e-80a4-0d9ff6f09459",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 框的类别\n",
    "bbox_class = {\n",
    "    'sjb_rect':0  \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7c896479-c286-41aa-bd52-4c93b8c0011c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 关键点的类别\n",
    "keypoint_class = ['angle_30', 'angle_60', 'angle_90']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b4fb7ee-0886-416c-bbaf-6c2dfe7e63c9",
   "metadata": {},
   "source": [
    "## 载入一个labelme格式的json标注文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "475e8269-be3e-48bf-b089-9538f63c4b19",
   "metadata": {},
   "outputs": [],
   "source": [
    "labelme_path = 'DSC_0281.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b17d0e7a-b987-4784-9dfe-a67ebfeefb15",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(labelme_path, 'r', encoding='utf-8') as f:\n",
    "    labelme = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2c2b711c-4e34-459b-854f-bfdda7170935",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['version', 'flags', 'shapes', 'imagePath', 'imageData', 'imageHeight', 'imageWidth'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979f8578-9c1c-4b32-b98e-53838d0ef5fe",
   "metadata": {},
   "source": [
    "### 元数据 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9fb38995-4ecc-44b4-a132-1d6851f18961",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'5.1.1'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['version']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "376d7a49-d5a1-4827-b68e-7334a231a144",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['flags']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1733f8f9-72d4-47f8-a010-2ab98b724afa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'DSC_0281.jpg'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像文件名\n",
    "labelme['imagePath']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c0411a3c-8e43-4878-8657-cbdca9e15c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "labelme['imageData']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9bc00c97-9d5f-4664-b502-631cdf2a15b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3712"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像高度\n",
    "labelme['imageHeight']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8185cf9e-84ba-4081-8e4a-db189b10d29e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5568"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像宽度\n",
    "labelme['imageWidth']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68bd2fd7-914f-4ebc-a65e-bf43ceb87761",
   "metadata": {},
   "source": [
    "### 标注信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2ddb2f62-b076-491f-83cf-aa4f12c82ea3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': 'sjb_rect',\n",
       "  'points': [[1138.0, 0.0], [4480.0, 1562.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}},\n",
       " {'label': 'sjb_rect',\n",
       "  'points': [[1160.0, 1954.0], [4560.0, 3566.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_30',\n",
       "  'points': [[4434.0, 284.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_30',\n",
       "  'points': [[4526.0, 2316.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_60',\n",
       "  'points': [[1202.0, 44.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_60',\n",
       "  'points': [[1218.0, 2000.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_90',\n",
       "  'points': [[1934.0, 1512.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'angle_90',\n",
       "  'points': [[1906.0, 3532.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'point',\n",
       "  'flags': {}},\n",
       " {'label': 'sjb_poly',\n",
       "  'points': [[1202.0, 28.0], [4454.0, 280.0], [1938.0, 1524.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'polygon',\n",
       "  'flags': {}},\n",
       " {'label': 'sjb_poly',\n",
       "  'points': [[1222.0, 1992.0], [4542.0, 2292.0], [1918.0, 3532.0]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'polygon',\n",
       "  'flags': {}}]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['shapes']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afca9bb8-8197-43d2-82b6-9a7e2b6f8dc2",
   "metadata": {},
   "source": [
    "## 生成YOLO格式的标注文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6485abbd-cd71-4704-bbd0-332320265e0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_width = labelme['imageWidth']   # 图像宽度\n",
    "img_height = labelme['imageHeight'] # 图像高度\n",
    "\n",
    "# 生成 YOLO 格式的 txt 文件\n",
    "suffix = labelme_path.split('.')[-2]\n",
    "yolo_txt_path = suffix + '.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca83553-c0f2-4a50-a6e1-79217a94fe87",
   "metadata": {},
   "source": [
    "## 遍历每个标注，如果遇到框，就找到该框里所有的关键点，并按顺序写入txt文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "72b09ebe-4971-4f6c-832e-73c6b12dff15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DSC_0281.json --> DSC_0281.txt 转换完成\n"
     ]
    }
   ],
   "source": [
    "with open(yolo_txt_path, 'w', encoding='utf-8') as f:\n",
    "    \n",
    "    for each_ann in labelme['shapes']: # 遍历每个标注\n",
    "\n",
    "        if each_ann['shape_type'] == 'rectangle': # 如果遇到框\n",
    "            \n",
    "            yolo_str = ''\n",
    "            \n",
    "            ## 框的信息\n",
    "            # 框的类别 ID\n",
    "            bbox_class_id = bbox_class[each_ann['label']]\n",
    "            yolo_str += '{} '.format(bbox_class_id)\n",
    "            # 左上角和右下角的 XY 像素坐标\n",
    "            bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "            bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "            bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "            bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "            # 框中心点的 XY 像素坐标\n",
    "            bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)\n",
    "            bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)\n",
    "            # 框宽度\n",
    "            bbox_width = bbox_bottom_right_x - bbox_top_left_x\n",
    "            # 框高度\n",
    "            bbox_height = bbox_bottom_right_y - bbox_top_left_y\n",
    "            # 框中心点归一化坐标\n",
    "            bbox_center_x_norm = bbox_center_x / img_width\n",
    "            bbox_center_y_norm = bbox_center_y / img_height\n",
    "            # 框归一化宽度\n",
    "            bbox_width_norm = bbox_width / img_width\n",
    "            # 框归一化高度\n",
    "            bbox_height_norm = bbox_height / img_height\n",
    "            \n",
    "            yolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm, bbox_width_norm, bbox_height_norm)\n",
    "            \n",
    "            ## 找到该框中所有关键点，存在字典 bbox_keypoints_dict 中\n",
    "            bbox_keypoints_dict = {}\n",
    "            for each_ann in labelme['shapes']: # 遍历所有标注\n",
    "                if each_ann['shape_type'] == 'point': # 筛选出关键点标注\n",
    "                    # 关键点XY坐标、类别\n",
    "                    x = int(each_ann['points'][0][0])\n",
    "                    y = int(each_ann['points'][0][1])\n",
    "                    label = each_ann['label']\n",
    "                    if (x>bbox_top_left_x) & (x<bbox_bottom_right_x) & (y<bbox_bottom_right_y) & (y>bbox_top_left_y): # 筛选出在该个体框中的关键点\n",
    "                        bbox_keypoints_dict[label] = [x, y]\n",
    "            \n",
    "            ## 把关键点按顺序排好\n",
    "            for each_class in keypoint_class: # 遍历每一类关键点\n",
    "                if each_class in bbox_keypoints_dict:\n",
    "                    keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_width\n",
    "                    keypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_height\n",
    "                    yolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm, 2) # 2-可见不遮挡 1-遮挡 0-没有点\n",
    "                else: # 不存在的点，一律为0\n",
    "                    yolo_str += '0 0 0 '.format(keypoint_x_norm, keypoint_y_norm, 0)\n",
    "            # 写入 txt 文件中\n",
    "            f.write(yolo_str + '\\n')\n",
    "            \n",
    "print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb08428-e4e8-4b4a-871e-0884cf004d00",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2460a3cd-fc21-47cc-9a0e-23db6bd9dc25",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
